import copy
import os
import numpy as np
# import torch
import json
import argparse
import mmap
import random
import matplotlib.pyplot as plt
from multiprocessing.pool import ThreadPool as Pool
# from torch.utils.data import Dataset, ConcatDataset
from mindspore.mindrecord import FileWriter
from mindformers.tools import logger
from internlm2_tokenizer import InternLM2Tokenizer

# http://spring.sensetime.com/pypi/packages/petrel_oss_sdk-2.2.0_19_gac5c680_master-py3-none-any.whl

IGNORE_TOKEN_ID = -100  # 补齐值设置为 -100

# 制作tokenize后的数据集
def transformer(sources, tokenizer, seq_length):
    """From alpaca to mindrecord."""
    input_ids = []
    labels = []
    special_tokenize = {"<s>": tokenizer.encode('')[0],  # bos_token
                        "</s>": tokenizer.encode('</s>', add_special_tokens=False)[0],  # eod_token
                        "newline": tokenizer.encode('\n', add_special_tokens=False)[0]}
    start_usr_token = tokenizer.encode('<|im_start|>user', add_special_tokens=False) + [special_tokenize["newline"]]
    start_ass_token = tokenizer.encode('<|im_start|>assistant', add_special_tokens=False) \
                      + [special_tokenize["newline"]]
    end_token = tokenizer.encode('<|im_end|>', add_special_tokens=False) + [special_tokenize["newline"]]
    len_sources = len(sources)
    index = 1
    for source in sources: # 遍历json对象
        if index % 10000 == 0:
            print(f"tokenize json对象({index}/{len_sources})")
        index += 1
        ins_token = tokenizer.encode(source['instruction'], add_special_tokens=False)
        inp_token = tokenizer.encode(source['input'], add_special_tokens=False) + [special_tokenize["newline"]]
        out_token = tokenizer.encode(source['output'], add_special_tokens=False)
# ----------------------------------------------------------------------------------------------------------------------------------------------------------------------
        # 1 构建input_id
        input_id = [special_tokenize["<s>"]] + [*start_usr_token, *ins_token, *end_token, *start_ass_token] + [*out_token, *end_token] # 去掉</s> eod_token
        # 2 构建label
        # label里的bos_token就是-100 是否保留 *end_token ？ 保留。
        label = [*out_token, *end_token] + [special_tokenize["</s>"]]
        # 在label前方进行pad -100 使label和input_id长度相等
        # 计算需要补齐的长度
        len_input_id = len(input_id)
        len_label = len(label)
        len_pad: int =  len_input_id - len_label
        # 使用 np.pad 在 `label` 前方补齐 `-100`
        label = np.pad(
            label,  # 需要补齐的数组
            (len_pad, 0),  # (前方补齐的长度, 后方补齐的长度)
            'constant',  # 使用常量值补齐
            constant_values=(IGNORE_TOKEN_ID, IGNORE_TOKEN_ID)  # 补齐值设置为 -100
        )

        input_ids.append(np.array(input_id).astype(np.int32))
        labels.append(np.array(label).astype(np.int32))

    return dict(input_ids=input_ids, labels=labels)

def read_json(json_path):
    """
    读取 JSON 文件并将其解析为字典列表。

    :param json_path: str, JSON 文件路径
    :return: list[dict], 包含字典的列表
    :raises: FileNotFoundError, JSONDecodeError
    """
    try:
        with open(json_path, 'r', encoding='utf-8') as f:
            metas = json.load(f)

        if isinstance(metas, list) and all(isinstance(item, dict) for item in metas):
            return metas
        else:
            raise ValueError("JSON 文件内容必须是字典列表。")
    except FileNotFoundError:
        raise FileNotFoundError(f"文件 {json_path} 不存在。")
    except json.JSONDecodeError as e:
        raise ValueError(f"JSON 文件格式错误: {e}")


# 随机化分组 考虑按照数据长度分组？
def random_group(token_lengths, seed, packed_length):
    rng = np.random.RandomState(seed)
    index = list(range(len(token_lengths)))
    rng.shuffle(index)  # 随机化索引

    pack_group = []
    token_length_sum = 0
    each_group = []
    for idx, sample_id in enumerate(index):
        sample_length = token_lengths[sample_id][1]
        if sample_length > packed_length:  # 如果单个样本的长度超过了4096，跳过该样本。
            continue
        token_length_sum += sample_length  # 更新总长度
        if token_length_sum >= packed_length:  # 如果总长度大于等于了4096，
            each_group.append(token_lengths[sample_id])
            pack_group.append(each_group)
            each_group = []
            token_length_sum = 0
        else:
            each_group.append(token_lengths[sample_id])
        # 这里注释掉 最后一个each_group除非正好到长度到达packed_length就结束 or 丢弃
        # if idx == len(token_lengths) - 1:  # 遍历到最后一个元组
        #     if len(each_group) > 0:
        #         pack_group.append(each_group)
    return pack_group

# 按照数据长度分组 待测试
def sort_group(token_lengths, packed_length):
    # 若不使用随机化索引， 对token_length按照length进行从小到大排序 使背包利用率最大化
    token_lengths = dict(sorted(token_lengths.items(), key=lambda item: item[1][1]))

    pack_group = []
    token_length_sum = 0
    each_group = []
    # v是元组
    for k,v in token_lengths.items():
        sample_length = v[1]
        if sample_length > packed_length:  # 如果单个样本的长度超过了4096，跳过该样本。
            continue
        token_length_sum += sample_length  # 更新总长度
        if token_length_sum >= packed_length:  # 如果总长度大于等于了4096，
            each_group.append(v)
            pack_group.append(each_group)
            each_group = []
            token_length_sum = 0
        else:
            each_group.append(v)
    return pack_group


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--mindrecord_schema", type=str, default="internlm2_alpaca")
    parser.add_argument("--input_glob", type=str, default="./merged_data_internlm2.json")  #JSON 文件路径
    parser.add_argument("--output_file", type=str, default="./merged_data_internlm2_packed.mindrecord")
    parser.add_argument("--model_file", type=str, default="./tokenizer.model")
    parser.add_argument("--file_partition", type=int, default=1)
    parser.add_argument("--seq_length", type=int, default=4096)
    args = parser.parse_args()

    print("读取原始json...")
    metas = read_json(args.input_glob)

    seq_length = args.seq_length
    seed = 1024
    tokenizer = InternLM2Tokenizer(vocab_file=args.model_file)
    eod_token = tokenizer.encode('</s>', add_special_tokens=False)[0]  # eod_token
    print("对json进行tokenize...")
    # 对metas进行tokenize 实际上metas就是alpaca mindrecord数据处理逻辑preprocess()的返回值
    metas = transformer(metas, tokenizer, seq_length)
    len_data = len(metas['input_ids'])  # 原始样本个数
    # print(f"{metas}")  # no bug
    print(f"len_data: {len(metas['input_ids'])}")  # no bug


    #metas结构
    #{
    #   'input_ids' = [[],[],[],[]...,[]],   这里相同索引的input_ids和labels的子列表长度相等
    #   'labels' = [[],[],[],[]...,[]]
    # }
    print("构造token_lengths...")
    lengths_dict = dict()
    # 将idx对应的token dict 化为((out_idx, length)元组 存入lengths_dict
    for idx in range(len_data):
        lengths_dict[idx] = (idx, len(metas['input_ids'][idx]))
    token_lengths = lengths_dict
    # test
    # print(f"token_lengths: {token_lengths}")  # no bug
    # token_lengths
    # {
    #     idx1: (idx1, length1), 这里是只存单个input_id的长度
    #     idx2: (idx2, length2),
    #     ...
    # }
    np.set_printoptions(threshold=np.inf)
    print("制作 pack_group...")
    # 使用随机化分组
    pack_group = random_group(token_lengths, seed, seq_length)
    # 或者使用按数据长度分组
    # pack_group = sort_group(token_lengths, seq_length)

    # print(f"pack_group :{pack_group}")
    print("-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------")
    # test
    # print(f"pack_group: {pack_group}")
    # 遍历每个子列表并计算每个子列表中所有元组的第二个元素之和
    def print_pack_group(pack_group):
        result = []
        for sublist in pack_group:
            sublist_sum = sum(t[1] for t in sublist)  # 使用列表生成式提取元组[1]并求和
            result.append(sublist_sum)
        # 打印结果
        for i, total in enumerate(result):
            print(f"pack_group 子列表 {i} 的所有元组 [1] 的和为: {total}")
    print_pack_group(pack_group)

    # pack_group =
    # [
    #   [(96,52),(71,658),(120,372)...,(33,77)],
    #   [(),(),(),(),()...,(),()]
    # ]

    #制作最终数据data
    data = dict()  # 为了数据集纠错 保存可视化格式数据集
    input_ids = []
    labels = []
    loss_masks = []
    position_ids = []
    attention_masks = []
    draw_cnt = 0
    len_pack_group = len(pack_group)
    # 遍历pack_group 制作最终数据data
    for idx, group in enumerate(pack_group):
        if idx % 100 == 0:
            print(f"遍历pack_group 制作最终数据data 进度: ({idx+1}/{len_pack_group})")
        input_id = []
        label = []
        loss_mask = []
        position_id = []
        attention_mask = []
        cnt = 1 # 用于制作 attention_mask
        # 按照group pack 一个batch
        for g in group:  # g是(index, length)的元组 length是index对应的input_id token的长度
            index, length = g
            new_length = len(metas["input_ids"][index])
            if new_length != length:  # 检测从pack_group取得的此input_id token的长度 和 测量的长度 是否相等
                print(index, f"current length {new_length} != cache {length}")
                continue
            # 1 构建input_id
            input_id.append(metas['input_ids'][index])
            # 2 构建label
            label.append(metas['labels'][index])

            len_input_id = len(metas['input_ids'][index])
            # 3 构建loss_mask
            len_pad = leading_100_count = np.argmax(metas['labels'][index] != -100)  # 计算label前导-100的数量
            len_old_label = len_input_id - len_pad
            loss_mask.append(np.concatenate([np.zeros(len_pad, dtype=np.int32), np.ones(len_old_label, dtype=np.int32)])) #  将此sample的loss_mask追加
            # 4 构建position_id
            position_id.append(np.arange(len_input_id, dtype=np.int32)) #  将此sample的position_id追加
            # 5 构建attention_mask no bug (1,1,1,2,2,2,2,2,3,3,3...)
            attention_mask.append(np.full(len_input_id, cnt, dtype=np.int32)) #  将此sample的attention_mask追加
            cnt += 1

        # test
        # print(f"input_id: {input_id}")
        # print(f"{len(input_id)}")
        # 合并为一维
        input_id = np.concatenate(input_id)  # no bug
        len_sample = len(input_id)  # 记录未截取的sample长度
        label = np.concatenate(label)# no bug
        loss_mask = np.concatenate(loss_mask)# no bug
        position_id = np.concatenate(position_id)# no bug
        attention_mask = np.concatenate(attention_mask)# no bug

        # if len_sample >= seq_length:
        # 截取至seq_length长度 无需补全eod
        input_id = input_id[:seq_length]
        label = label[:seq_length]
        loss_mask = loss_mask[:seq_length]
        position_id = position_id[:seq_length]
        attention_mask = attention_mask[:seq_length]
        # elif idx == len(pack_group) - 1:  # 若为最后一个group 前面逻辑将最后一个group丢弃了 所以这里不会走
        #     pad_len = seq_length - len_sample
        #     input_id = np.pad(input_id, (0, pad_len), mode='constant', constant_values=0)
        #     label = np.pad(label, (0, pad_len), mode='constant', constant_values=IGNORE_TOKEN_ID)
        #     # loss_mask = loss_mask[:seq_length]       # 若不丢弃最后一个group 待修改
        #     # position_id = position_id[:seq_length]   # 若不丢弃最后一个group 待修改
        # else:
        #     raise Exception("pack_group未遍历到末尾group 就出现len_sample < seq_length")

        # 追加
        np.set_printoptions(linewidth = 5000)
        input_ids.append(input_id)
        data[f"      input_id{idx}"] = input_id  # 记录
        labels.append(label)
        data[f"         label{idx}"] = label  # 记录
        loss_masks.append(loss_mask)
        data[f"     loss_mask{idx}"] = loss_mask  # 记录
        position_ids.append(position_id)
        data[f"   position_id{idx}"] = position_id  # 记录
        attention_masks.append(attention_mask)
        data[f"attention_mask{idx}"] = attention_mask  # 记录
    # for循环结束
    # check

    import json
    import numpy as np
    def align_and_save_dict_to_jsonlines(data, file_path):
        """
        将字典数据保存为对齐后的 JSON Lines 格式（每行一个 JSON 对象，写入文件时连续写入）。数字对齐后，可读性更好
        支持 numpy 数组类型的数据自动转换为列表。

        :param data: dict | 字典数据。
        :param file_path: str | 保存的文件路径。
        """
        def convert_numpy(obj):
            if isinstance(obj, np.ndarray):
                return obj.tolist()
            raise TypeError(f"Object of type {obj.__class__.__name__} is not JSON serializable")
        # 检查数据是否为字典
        if not isinstance(data, dict):
            raise ValueError("数据应为字典类型")
        # 确定数值的最大长度
        max_lengths = {}
        for key, values in data.items():
            if isinstance(values, list):
                max_lengths[key] = max(len(str(v)) for v in values)
        print("打开文件以写入 JSON Lines 格式数据")
        with open(file_path, 'w', encoding='utf-8') as f:
            for key, value in data.items():
                # 对齐数值列表
                aligned_values = ','.join([f"{str(v):>6}" for v in value])
                json_line = {"key": key, "value": aligned_values}
                json.dump(json_line, f, ensure_ascii=False, default=convert_numpy)
                f.write('\n')  # JSON Lines 每条记录以换行符分隔
        print(f"数据已成功保存为对齐后的 JSON Lines 格式到 {file_path}")

    # align_and_save_dict_to_jsonlines(data, "/data/version/MS24_1022/data/file/dict.jsonl")  # 保存可视化数据json 用于检查数据
    print(f"len input_ids: {len(input_ids)}")
    print(f"len labels: {len(labels)}")
    print(f"len loss_masks: {len(loss_masks)}")
    print(f"len position_ids: {len(position_ids)}")
    print(f"len attention_masks: {len(attention_masks)}")
    print("将数据写入mindrecord文件...")
    #将数据写入mindrecord文件
    out_dir, out_file = os.path.split(os.path.abspath(args.output_file))
    if not os.path.exists(out_dir):
        os.mkdir(out_dir)
    print("准备数据模式 (schema)...")
    schema = {'input_ids': {"type": "int32", "shape": [-1]},  # ms规定数据类型
              'labels': {"type": "int32", "shape": [-1]},
              'loss_masks': {"type": "int32", "shape": [-1]},
              'position_ids': {"type": "int32", "shape": [-1]},
              'attention_masks': {"type": "int32", "shape": [-1]}}  #

    print("初始化写入器（FileWriter）...")
    writer = FileWriter(file_name=args.output_file,
                        shard_num=args.file_partition)
    writer.set_page_size(256 * 1024 * 1024)  # 设置页面大小为 256MB，根据seq_length调整 范围[32768 bytes, 268435456 bytes] 单条数据最大8K
    writer.add_schema(schema, args.mindrecord_schema)

    len_packed = len(input_ids)  # packed batch数量
    print(f"len_packed: {len_packed}")
    # 都是长度相同的列表或 numpy 数组
    for i in range(len_packed):  # 遍历数据
        # 为每个样本准备数据
        x = {
            'input_ids': input_ids[i],
            'labels': labels[i],
            'loss_masks': loss_masks[i],
            'position_ids': position_ids[i],
            'attention_masks': attention_masks[i]
            }
        # print(f"eod_token: {eod_token}")
        print(f"shape 第{i+1}条x:")
        print(f"shape input_id: {input_ids[i].shape}")
        # print(input_ids[i][4095])
        print(f"shape label: {labels[i].shape}")
        # print(labels[i][4095])
        print(f"shape loss_mask: {loss_masks[i].shape}")
        print(f"shape position_id: {position_ids[i].shape}")
        print(f"shape attention_mask: {attention_masks[i].shape}")
        # exit()
        # 将当前数据写入文件
        print(f"写入x({i+1}/{len_packed})...")
        writer.write_raw_data([x])  # write_raw_data 接受一个字典列表
    # 提交并完成文件写入
    writer.commit()
    print(f"已转化并保存 {len(input_ids)} 条记录。")

# alpaca_data_preprocess.py 中打印x:
# <class 'dict'> {'input_ids': array([    1, 92543,  1008,   364, 35126,  2478, 10552,   500, 19597,
#         9444,   281, 92542,   364, 92543,   525, 11353,   364], dtype=int32),
#         'labels': array([  312,  5302,   397,   395, 24198, 10105,   454,  1426,  2866,
#          442,  3089, 11411,   446, 25757,   454, 24246,   281,   848,
#          314,   281, 32935, 15630,   442,  2662,   829,  2642,  4703,
#          454,  3906,   281,   848,   308,   281,  2280,  3485,  6237,
#          454, 10442,   395, 13116,  6237,  9836,   281, 92542,   364,   2], dtype=int32)}
























